package controllers

import (
	"fmt"
	kaiv1alpha1 "github.com/AliyunContainerService/et-operator/api/v1alpha1"
	logger "github.com/sirupsen/logrus"
	corev1 "k8s.io/api/core/v1"
	"k8s.io/apimachinery/pkg/api/resource"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"path"
	"strings"
)

func (r *TrainingJobReconciler) newLauncherPod(job *kaiv1alpha1.TrainingJob) *corev1.Pod {
	launcherName := job.Name + launcherSuffix
	labels := GenLabels(job.Name)
	labels[labelTrainingRoleType] = launcher
	podSpec := job.Spec.ETReplicaSpecs.Launcher.Template.DeepCopy()
	// copy the labels and annotations to pod from PodTemplate
	if len(podSpec.Labels) == 0 {
		podSpec.Labels = make(map[string]string)
	}
	for key, value := range labels {
		podSpec.Labels[key] = value
	}
	podSpec.Spec.InitContainers = append(podSpec.Spec.InitContainers, r.newLauncherInitContainer(job))
	//podSpec.Spec.InitContainers = append(podSpec.Spec.InitContainers, kubedeliveryContainer())
	if len(podSpec.Spec.Containers) == 0 {
		logger.Errorln("Launcher pod does not have any containers in its spec")
		return nil
	}

	container := podSpec.Spec.Containers[0]
	container.VolumeMounts = append(container.VolumeMounts,
		corev1.VolumeMount{
			Name:      hostfileVolumeName,
			MountPath: hostfileMountPath,
		},
		corev1.VolumeMount{
			Name:      configVolumeName,
			MountPath: configMountPath,
		},
		corev1.VolumeMount{
			Name:      kubectlVolumeName,
			MountPath: kubectlMountPath,
		},
		corev1.VolumeMount{
			Name:      deepSpeedHostfileName,
			MountPath: deepSpeedMountPath,
		})

	if job.GetAttachMode() == kaiv1alpha1.AttachModeKubexec {
		container.Env = append(container.Env, corev1.EnvVar{
			Name:  "OMPI_MCA_plm_rsh_agent",
			Value: getKubexecPath(),
		})
	}
	podSpec.Spec.Containers[0] = container
	podSpec.Spec.ServiceAccountName = launcherName

	setRestartPolicy(podSpec)
	hostfileMode := int32(0444)
	scriptMode := int32(0555)

	podSpec.Spec.Volumes = append(podSpec.Spec.Volumes,
		corev1.Volume{
			Name: hostfileVolumeName,
			VolumeSource: corev1.VolumeSource{
				EmptyDir: &corev1.EmptyDirVolumeSource{},
			},
		},
		corev1.Volume{
			Name: deepSpeedHostfileName,
			VolumeSource: corev1.VolumeSource{
				EmptyDir: &corev1.EmptyDirVolumeSource{},
			},
		},
		corev1.Volume{
			Name: kubectlVolumeName,
			VolumeSource: corev1.VolumeSource{
				EmptyDir: &corev1.EmptyDirVolumeSource{},
			},
		},
		corev1.Volume{
			Name: configVolumeName,
			VolumeSource: corev1.VolumeSource{
				ConfigMap: &corev1.ConfigMapVolumeSource{
					LocalObjectReference: corev1.LocalObjectReference{
						Name: job.Name + configSuffix,
					},
					Items: []corev1.KeyToPath{
						{
							Key:  hostfileName,
							Path: hostfileName,
							Mode: &hostfileMode,
						},
						{
							Key:  discoverHostName,
							Path: discoverHostName,
							Mode: &hostfileMode,
						},
						{
							Key:  kubexeclFileName,
							Path: kubexeclFileName,
							Mode: &scriptMode,
						},
						{
							Key:  deepSpeedHostfileName,
							Path: deepSpeedHostfileName,
							Mode: &hostfileMode,
						},
					},
				},
			},
		})
	return &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:        launcherName,
			Namespace:   job.Namespace,
			Labels:      podSpec.Labels,
			Annotations: podSpec.Annotations,
			OwnerReferences: []metav1.OwnerReference{
				*metav1.NewControllerRef(job, kaiv1alpha1.SchemeGroupVersionKind),
			},
		},
		Spec: podSpec.Spec,
	}
}

func (r *TrainingJobReconciler) newLauncherInitContainer(job *kaiv1alpha1.TrainingJob) corev1.Container {
	originHostfilePath := path.Join(configMountPath, hostfileName)
	mountHostfilePath := getHostfilePath(job)
	cpHostfile := fmt.Sprintf("cp %s %s && chmod 600 %s",
		originHostfilePath,
		mountHostfilePath,
		mountHostfilePath)
	originDiscoverHostPath := path.Join(configMountPath, discoverHostName)
	discoverHostPath := path.Join(hostfileMountPath, discoverHostName)
	cpDiscoverHostfile := fmt.Sprintf("cp %s %s && chmod +x %s",
		originDiscoverHostPath,
		discoverHostPath,
		discoverHostPath)

	originDPHostfilePath := path.Join(configMountPath, deepSpeedHostfileName)
	dpHostPath := path.Join(deepSpeedMountPath, hostfileName)
	cpDPHostfile := fmt.Sprintf("cp %s %s && chmod 600 %s",
		originDPHostfilePath,
		dpHostPath,
		dpHostPath)

	return corev1.Container{
		Name:            initContainerName,
		Image:           r.InitContainerImage,
		ImagePullPolicy: corev1.PullIfNotPresent,
		VolumeMounts: []corev1.VolumeMount{
			{
				Name:      hostfileVolumeName,
				MountPath: hostfileMountPath,
			},
			{
				Name:      deepSpeedHostfileName,
				MountPath: deepSpeedMountPath,
			},
			{
				Name:      configVolumeName,
				MountPath: configMountPath,
			},
		},
		Command: []string{
			"sh",
			"-c",
			strings.Join([]string{cpHostfile, cpDiscoverHostfile, cpDPHostfile}, " && "),
		},
		Resources: corev1.ResourceRequirements{
			Limits: corev1.ResourceList{
				corev1.ResourceCPU:              resource.MustParse(initContainerCpu),
				corev1.ResourceMemory:           resource.MustParse(initContainerMem),
				corev1.ResourceEphemeralStorage: resource.MustParse(initContainerEphStorage),
			},
			Requests: corev1.ResourceList{
				corev1.ResourceCPU:              resource.MustParse(initContainerCpu),
				corev1.ResourceMemory:           resource.MustParse(initContainerMem),
				corev1.ResourceEphemeralStorage: resource.MustParse(initContainerEphStorage),
			},
		},
	}
}
